#include "StdAfx.h"
#include "PlayerTracker.h"

using namespace std;
using namespace cv;
//Simple KLT tracker using opencv's implementation.


PlayerTracker::PlayerTracker(){
}

//construct tracker and extract feature points in the given rectangles.
PlayerTracker::PlayerTracker(Mat startingFrame, Rect p1Win, Rect p2Win) : _frame(startingFrame), 
	_p1Win(p1Win), _p2Win(p2Win) {
	
	startingFrame.copyTo(_startFrame);
	cvtColor(_startFrame, _startGray, CV_BGR2GRAY);
	_frame = _startGray;

	_p1 = Player();
	_p2 = Player();

	_frameNum = 1;

	getFilteredTrackingPts(_startFrame, p1Win, _currP1Pts, p2Win, _currP2Pts);
	filterPointsBasedOnLocation();

	_p1.addPoints(_currP1Pts);
	_p2.addPoints(_currP2Pts);
}

void PlayerTracker::updateFrame(Mat newFrame) {

	_prevFrame = _frame;
	_prevP1Pts = _currP1Pts;
	_prevP2Pts = _currP2Pts;

	Mat newGray;
	cvtColor(newFrame, newGray, CV_BGR2GRAY);

	_frame = newGray;

	//Optical Flow Step ----------------------------------------------------------------------//
	performOpFlow(_prevFrame, _frame ,_prevP1Pts, _currP1Pts, _prevP2Pts, _currP2Pts);
	
	filterPointsBasedOnLocation();
	_p1.addPoints(_currP1Pts);
	_p2.addPoints(_currP2Pts);
	

	//Draws Op Flow result
	for(size_t c = 0; c < _currP1Pts.size(); c++) {
		circle(newFrame, _currP1Pts[c], 4, Scalar(0,255,0));
		line(newFrame, _prevP1Pts[c], _currP1Pts[c], Scalar(0,255,0));
	}

	for(size_t c = 0; c < _currP2Pts.size(); c++) {
		circle(newFrame, _currP2Pts[c], 4, Scalar(0,255,0));
		line(newFrame, _prevP2Pts[c], _currP2Pts[c], Scalar(0,255,0));

	}

	//Draw Centroids
	Point2f p1Centroid;
	_p1.centroid(p1Centroid);
	Point2f p2Centroid;
	_p2.centroid(p2Centroid);

	////Generate new feature points to track if current number of feature points are less than a certain number.
	//if(_currP1Pts.size()<5 ){
	//	Rect p1Rect(p1Centroid.x-10, p1Centroid.y-10, 20, 20);
	//	Rect p2Rect(p2Centroid.x-10, p2Centroid.y-10, 20, 20);
	//	_currP1Pts.clear();
	//	vector<cv::Point2f> tempPoints;
	//	getFilteredTrackingPts(newFrame, p1Rect, _currP1Pts, p2Rect, tempPoints);
	//	cout<< "Getting new points P1"<< endl;
	//}

	//if(_currP2Pts.size()<5 ){
	//	Rect p1Rect(p1Centroid.x-10, p1Centroid.y-10, 20, 20);
	//	Rect p2Rect(p2Centroid.x-10, p2Centroid.y-10, 20, 20);
	//	_currP2Pts.clear();
	//	vector<cv::Point2f> tempPoints;
	//	getFilteredTrackingPts(newFrame, p1Rect, tempPoints, p2Rect, _currP2Pts);
	//	cout<< "Getting new points P1"<< endl;
	//}


	circle(newFrame, p1Centroid, 4, Scalar(0,0,255), 4);
	circle(newFrame, p2Centroid, 4, Scalar(0,0,255), 4);
}

void PlayerTracker::filterPointsBasedOnLocation(){
	
	
	// Get centroid point
	double player1XMean, player2XMean, player1YMean, player2YMean;
	double xSum, ySum;
	xSum = ySum = 0;
	for(int i=0; i<_currP1Pts.size(); i++){
		xSum += _currP1Pts[i].x;
		ySum += _currP1Pts[i].y;
	}
	player1XMean = xSum/_currP1Pts.size(); 
	player1YMean = ySum/_currP1Pts.size(); 

	xSum = ySum = 0;
	for(int i=0; i<_currP2Pts.size(); i++){
		xSum += _currP2Pts[i].x;
		ySum += _currP2Pts[i].y;
	}
	player2XMean = xSum/_currP2Pts.size(); 
	player2YMean = ySum/_currP2Pts.size(); 


	//filter based on distance from centroid
	double p1Thresh, p2Thresh;
	double distanceSum =0;
	vector<double> distances;
	for(int i=0; i<_currP1Pts.size(); i++){
		double distance = sqrt( pow((_currP1Pts[i].x - player1XMean),2) + pow((_currP1Pts[i].y - player1YMean),2) );
		distances.push_back(distance);
		distanceSum += distance;
	}
	int p1Removed =0;
	p1Thresh = distanceSum/_currP1Pts.size(); 
	//adding some slack, otherwise points are always eliminated.
	p1Thresh *= 1.5;
	for(int i=0; i<_currP1Pts.size(); i++){
		if(distances[i] > p1Thresh && distances[i]>0.01*_frame.cols){
			_currP1Pts.erase(_currP1Pts.begin()+i);
			p1Removed ++ ;
		}
	}
	
	//removing player 1 feature points that are too far away from the mean.
	distanceSum =0;
	distances.clear();
	for(int i=0; i<_currP2Pts.size(); i++){
		double distance = sqrt( pow((_currP2Pts[i].x - player2XMean),2) + pow((_currP2Pts[i].y - player2YMean),2) );
		distances.push_back(distance);
		distanceSum += distance;
	}
	
	//removing player 1 feature points that are too far away from the mean.
	int p2Removed =0;
	p2Thresh = distanceSum/_currP2Pts.size();
	p2Thresh *= 2;
	for(int i=0; i<_currP2Pts.size(); i++){
		if(distances[i] > p2Thresh && distances[i]>0.01*_frame.cols){
			_currP2Pts.erase(_currP2Pts.begin()+i);
			p2Removed ++;
		}
	}

	if(p1Removed>0 || p2Removed>0){
		cout<<"Points Filter : p1 Removed: " << p1Removed << " p2Removed: " << p2Removed << endl;
		cout<<"p1 Size: " << _currP1Pts.size() << " p2 Size: " << _currP2Pts.size() << endl;
	}
}


void PlayerTracker::performOpFlow(Mat prevFrame, Mat newFrame, vector<Point2f> prevP1Pts, vector<Point2f> &newP1Pts, vector<Point2f> prevP2Pts, vector<Point2f> &newP2Pts) {
	vector<uchar> status;
    vector<float> err;

	vector<Point2f> points1;
	newP1Pts = vector<Point2f>();
	calcOpticalFlowPyrLK(prevFrame, newFrame, prevP1Pts, points1, status, err);

	int added = 0;
	for(size_t i = 0; i < points1.size(); i++) {
		if(status[i] == 1){
			Point p1 = points1[i];
			newP1Pts.push_back(p1);
			added++;
		}
	}
	int p1Removed = points1.size()-added;


	added =0;
	status.clear();
	err.clear();
	vector<Point2f> points2;
	newP2Pts = vector<Point2f>();
	calcOpticalFlowPyrLK(prevFrame, newFrame, prevP2Pts, points2, status, err);
	for(size_t i = 0; i < points2.size(); i++) {
		if(status[i] == 1){
			Point p2 = points2[i];
			newP2Pts.push_back(p2);
			added++;
		}
	}
	int p2Removed = points2.size()-added;

	if(p1Removed>0 || p2Removed>0){
		cout<<"Optical Flow : p1 Removed: " << p1Removed << " p2Removed: " << p2Removed << endl;
		cout<<"p1 Size: " << _currP1Pts.size() << " p2 Size: " << _currP2Pts.size() << endl;
	}

}

//extract feature point in the box around the player.
//video has different qualities, we try multiple times till get desecent number of feature points
void PlayerTracker::getFilteredTrackingPts(Mat mat, Rect w1, vector<Point2f> &tPtsP1, Rect w2, vector<Point2f> &tPtsP2) {
	tPtsP1.clear();
	tPtsP2.clear();

	Mat gray;
	cvtColor(mat, gray, CV_BGR2GRAY);
	vector<Point2f> allCorners;
	double quality = 0.01;
	goodFeaturesToTrack(gray, allCorners, 0, quality, 5);
	for(size_t i = 0; i < allCorners.size(); i++) {
		Point2f pt = allCorners[i];

		if(w1.contains(pt)) {
			tPtsP1.push_back(pt);
		}
		else if(w2.contains(pt)) {
			tPtsP2.push_back(pt);
		}
	}

	//if not enough feature points detected keep halving the quality of ht points.
	int tries = 8;
	while(tries > 0 && tPtsP1.size()<5 && tPtsP2.size()<5){
		quality = quality/2.0;
		goodFeaturesToTrack(gray, allCorners, 0, quality, 5);
		for(size_t i = 0; i < allCorners.size(); i++) {
			Point2f pt = allCorners[i];
			if(w1.contains(pt)) {
				tPtsP1.push_back(pt);
			}
			else if(w2.contains(pt)) {
				tPtsP2.push_back(pt);
			}
		}	
	}
}

Player PlayerTracker::getPlayer(int playerNum){
	if(playerNum ==1){
		return _p1;
	}
	else{
		return _p2;
	}
}

PlayerTracker::~PlayerTracker(void)
{
}
